#include "snp_cand.hpp"
#include "option_parse.hpp"
#include "io.hpp"
#include "util.hpp"

void do_mapcand(int argc, char **argv)
{
    prog_options options = parse_mapcand_options(argc, argv);

    map_str_vafl_t subs, totl;
    cand_blocks_t cand_blocks;
    cand_vector_t snp_cands;
    std::map<std::string, std::string> seqs;
    mapping_type t;

    // TODO: also read the mapping type. and support sequence space.
    read_map_count_file(options.map_count_file.c_str(), t, subs, totl);

    if(t & COLOR)
    {
        find_snp_cands(options, cand_blocks, snp_cands, subs, totl);
    }
    else
    {
        find_seq_snp_cands(options, cand_blocks, snp_cands, subs, totl);
    }

    parse_fasta(options.fasta_file.c_str(), seqs);

    attach_seqs_to_snp_cands(seqs, snp_cands);

    write_snp_cand(std::cout, cand_blocks, snp_cands);
}

void attach_seqs_to_snp_cands(std::map<std::string, std::string> &seqs, cand_vector_t &snp_cands)
{
    for(unsigned int i = 1; i < snp_cands.size(); i++)
        snp_cands[i].seq = seqs[snp_cands[i].chr].substr(snp_cands[i].start, snp_cands[i].end - snp_cands[i].start);
}

void do_mapcount(int argc, char **argv)
{
    prog_options options = parse_mapcount_options(argc, argv);
    
    /* load the lengths of all the genomes. */
    map_str_ui_t ref_seq_len = build_fasta_index(options.fasta_file);
    
    mapping_type t = get_mapping_type(options.map_file.c_str());

    map_str_vafl_t subs, totl;

    if(t & COLOR)
        find_subs_totl(options, subs, totl, ref_seq_len);
    else
        find_seq_subs_totl(options, subs, totl, ref_seq_len);


    write_map_count(std::cout, t, subs, totl);
}


void do_snpcand(int argc, char **argv)
{
    prog_options options = parse_snpcand_options(argc, argv);

    /* load the lengths of all the genomes. */
    map_str_ui_t ref_seq_len = build_fasta_index(options.fasta_file);

    mapping_type t = get_mapping_type(options.map_file.c_str());

    if(t & COLOR )
    {
        map_str_vafl_t subs, totl;
        cand_blocks_t cand_blocks;
        cand_vector_t snp_cands;

        find_subs_totl(options, subs, totl, ref_seq_len);
        find_snp_cands(options, cand_blocks, snp_cands, subs, totl);

        std::map<std::string, std::string> seqs;
        parse_fasta(options.fasta_file.c_str(), seqs);
        attach_seqs_to_snp_cands(seqs, snp_cands);

        write_snp_cand(std::cout, cand_blocks, snp_cands);
    }
    else
    {
        map_str_vafl_t subs, totl;
        cand_blocks_t cand_blocks;
        cand_vector_t snp_cands;

        std::map<std::string, std::string> seqs;
        find_seq_subs_totl(options, subs, totl, ref_seq_len);
        find_seq_snp_cands(options, cand_blocks, snp_cands, subs, totl);

        parse_fasta(options.fasta_file.c_str(), seqs);
        attach_seqs_to_snp_cands(seqs, snp_cands);

        write_snp_cand(std::cout, cand_blocks, snp_cands);
    }
}

/* 
  Decide if the two color sequence read is a valid SNP candidates for the given reference. 
 */
bool two_colors(const char *read, const char *ref)
{
    if( read[0] == read[1] && ref[0] == ref[1] )
        return true;

    if( read[0] == ref[1] && read[1] == ref[0])
        return true;

    if( !( read[0] == ref[0] || read[0] == ref[1]) && !(read[1] == ref[0] || read[1] == ref[1]) && ref[0] != ref[1] && read[0] != read[1] )
        return true;

    return false;
}

/* 
  Decide if the read contains valid snps for the sequence ref. 
 */
bool legal(const char *read, const char *ref, unsigned int len)
{
    /* single color changes are never valid snp candidates. */
    if( len == 1 ) return false;

    /* two color changes are special cased above. */
    if( len == 2 ) return two_colors(read, ref);

    /* Otherwise since the sequence of colors is independant of the starting
     * base, generate two trial sequences according to the colors given,
     * starting each with 'A', and see if they reach alignment before the end
     * of the sequence. 
     **/
    char re = 'A';
    char rf = 'A';

    //assert( ! (ref[0] == read[0] || ref[len-1] == read[len-1]) );

    for(unsigned int i = 0; i < len; i++)
    {
        re = next_color[re][read[i]];
        rf = next_color[rf][ref[i]];
    }
    
    return (re == rf);
}


void find_subs_totl(prog_options &options, map_str_vafl_t &subs, map_str_vafl_t &totl, map_str_ui_t &ref_seq_len)
{

    /* calculate the read length. */
    // TODO: add error handling if problem with file.
    unsigned int read_length = get_read_length(options.map_file.c_str());

    /* number of total reads that map to each location */
    //std::map<std::string, std::valarray<float> > totl;
    /* number of errors in reads at each site. */
    //std::map<std::string, std::valarray<float> > subs;

    std::map<std::string, unsigned int>::iterator mi;

    /* resize all of the arrays to be the size of the genomes found in the index. */
    for(mi = ref_seq_len.begin(); mi != ref_seq_len.end(); ++mi)
    {
        subs[mi->first].resize(mi->second);
        totl[mi->first].resize(mi->second);
    }

    ColorMappingParser parser(options.map_file.c_str());

    std::vector<float> read_snp_cand;
    read_snp_cand.resize(read_length);

    int total                         = 0;
    int skipped_cutoff                = 0;
    int skipped_unique                = 0;
    

    /* Iterate over the mapping file. */
    ColorRead r;
    while(parser.next_read(r))
    {
        total++;

        if( options.unique && r.map_locs.size() > 1)
        {
            skipped_unique++;
            continue;
        }

        if( r.map_locs.size() > options.max_map_loci )
        {
            skipped_cutoff++;
            continue;
        }

        for(unsigned int j = 0; j < r.map_locs.size(); j++)
        {
            MapLoci m = r.map_locs[j];

            //std::cout << "m.id = " << m.id << " m.pos = " << m.pos << std::endl;

            if(totl.find(m.id) == totl.end())
            {
                if(options.ignore_missing)
                    continue;
                else
                    error("chromosome " + m.id + " does not exist in the fasta index file. (did you forget --ignore_missing?)");
            }

            for(unsigned int k = 0; k < read_length; k++)
                read_snp_cand[k] = 0.0;
            

            // add read to the totl array.
            for(unsigned int i = 0; i < r.seq.size(); i++)
                    totl[m.id][m.pos+i] += 1.0 / r.map_locs.size();

            // find which mismatches are valid snp mismatches.
            for(unsigned int i = 1; i < r.seq.size() - 1; i++)
            {
                // if we get a match then a mismatch, consider this a candidate.
                if(r.seq[i] == m.seq[i]  && r.seq[i+1] != m.seq[i+1])
                {
                    // look as far as we allow a consecutive snp block.
                    for(unsigned int k = i+2; k < i + 3 + options.consec_snps; k++)
                    {
                        // TODO: move this into the for loop condition.
                        // do not allow the candidate to run off the end of the sequence.
                        if( k + 1 >= r.seq.size() )
                            break;

                        // if the sequence gets back on track (legal) we add
                        // all the parts as snp candidates. 
                        if(r.seq[k] != m.seq[k] && r.seq[k+1] == m.seq[k+1])
                        {
                            if(legal(&r.seq.c_str()[i+1], &m.seq.c_str()[i+1], k - i))
                            {
                                for(unsigned int x = i+1; x < k; x++)
                                    read_snp_cand[x] = 1.0;
                                
                                // we start over at the next adjacent position when scanning.
                                i = i + 1;
                                break;
                            }
                        }
                    }
                }
            }

            // mark the substitution locations on the genome. 
            // ambig_partial == True
            if( m.strand == '+' )
                for(unsigned int k = 0; k < read_length; k++)
                    subs[m.id][m.pos + k] += read_snp_cand[k] / r.map_locs.size(); 
            else // m.strand == '-'
                for(unsigned int k = 0; k < read_length; k++)
                    subs[m.id][m.pos + r.seq.size() - k - 1] += read_snp_cand[k] / r.map_locs.size(); 
        }
    }
}

void find_snp_cands(prog_options &options, cand_blocks_t &cand_blocks, cand_vector_t &snps, map_str_vafl_t &subs, map_str_vafl_t &totl)
{

    // We store the SNPs in two ways for fast lookup.  An array the size of the
    // genome, so we can quickly determine with a lookup whether the size is
    // covered with a SNP.  And a list of each snp block, for storing the information about the SNPs.

    std::map<std::string, std::valarray<float> >::iterator mi;

    // generate candidates from subs/totl.
    
    /* Our final set of snp candidates to return. */
    std::map<std::string, std::valarray<bool> > snp_cand;
    
    for(mi = subs.begin(); mi != subs.end(); ++mi)
    {
        std::string  s = mi->first;

        snp_cand[s].resize(mi->second.size());
        snp_cand[s] = false;

        for(unsigned int cti = 0; cti < options.cover_threshold.size(); cti++)
        {
            int coverage  = options.cover_threshold[cti].first;
            float threshold = options.cover_threshold[cti].second;

            snp_cand[s] |= (totl[s] >= (float)coverage) && (subs[s] / totl[s] > threshold);
        }

    }
        
    // push dummy block into pos = 0, so that we do not have to do -1
    // adjustment later to access the snps in the array.
    std::string fake_chr("FAKE SNP");
    std::valarray<float> _subs, _totl;
    //std::vector<double> _subs, _totl;
    snps.push_back(SnpRecord( -1, fake_chr, -1, -1, _subs, _totl ) );
 
    // TODO: is there a better way to do this with valarray.cshift()?

    unsigned int snp_id = 1;

    // Find beginning and end of snp blocks, then walk over sequence to mark blocks.
    for(mi = subs.begin(); mi != subs.end(); ++mi)
    {
        std::string chr = mi->first;
        unsigned int len = mi->second.size();

        std::valarray<float> diff(0.0, len + 2);

        diff[0]   = snp_cand[chr][0] - 0;
        diff[len] = 0 - snp_cand[chr][len-1];
        
        for(unsigned int i = 1; i < len; i++)
            diff[i] = snp_cand[chr][i] - snp_cand[chr][i-1];

        std::valarray<bool> sblock = (diff > (float)0.0);
        std::valarray<bool> eblock = (diff < (float)0.0);

        cand_blocks[chr].resize(len, 0);

        bool in_block = false;
        
        unsigned int start=0, end=0;

        for( unsigned int i = 0; i < len+2; i++ )
        {
            if( sblock[i] )
            {
                start = i;
                in_block = true;
            }
            if( eblock[i] )
            {
                assert(in_block);
                in_block = false;
                end = i;

                unsigned int j, k;
                
                if(end - start <= options.max_block_size)
                {
                    _subs.resize(end-start); _totl.resize(end-start);
                    _subs = subs[chr][std::slice(start, end - start, 1)];
                    _totl = totl[chr][std::slice(start, end - start, 1)];

                    snps.push_back(SnpRecord( snp_id, chr, start, end, _subs, _totl) );
                    for(k = start; k < end; k++)
                        cand_blocks[chr][k] = snp_id;
                    snp_id++;
                }
                /* Here is where we split long snp blocks into shorter ones. */
                else
                {
                    // fliter blocks
                    std::valarray<float> local_block = (std::valarray<float>)subs[chr][std::slice(start, end-start, 1)] / (std::valarray<float>)totl[chr][std::slice(start, end-start, 1)];
                    std::valarray<bool> status(false, end - start);
k=0;
                    while(k+options.max_block_size< end-start)
                    {
                        int idx= std::min_element(&local_block[k],&local_block[k+options.max_block_size+1]) - &local_block[0];
                        status[idx]=true;
                        k=idx+1;
                    }

                    unsigned int r = 0;
                    while(r < status.size())
                    {
                        while(status[r] && r < status.size())
                            r++;
                        j = r;
                        while(!status[j] && j < status.size())
                            j++;
                        
                        _subs.resize(j-r); _totl.resize(j-r);
                        _subs = subs[chr][std::slice(r + start, j - r, 1)];
                        _totl = totl[chr][std::slice(r + start, j - r, 1)];
                        
                        snps.push_back(SnpRecord(snp_id, chr, r+start, j+start, _subs, _totl));

                        for(k = start + r; k < start + j; k++)
                            cand_blocks[chr][k] = snp_id;

                        r = j+1;
                        snp_id++;
                    }
                }
            }
        }

    }
    if(snp_id == 1)
    { 
        std::cerr << "comb snpcand did not generate any snp candidates!" << std::endl;
        std::cerr << "Consider modifying the cover/threshold parameters." << std::endl;
    }
}

/*******************************************************
 * SEQUENCE FUNCTIONS.
*/
void find_seq_subs_totl(prog_options &options, map_str_vafl_t &subs, map_str_vafl_t &totl, map_str_ui_t &ref_seq_len)
{

    /* calculate the read length. */
    // TODO: add error handling if problem with file.
    unsigned int read_length = get_read_length(options.map_file.c_str());

    std::map<std::string, unsigned int>::iterator mi;

    /* resize all of the arrays to be the size of the genomes found in the index. */
    for(mi = ref_seq_len.begin(); mi != ref_seq_len.end(); ++mi)
    {
        subs[mi->first].resize(mi->second);
        totl[mi->first].resize(mi->second);
    }

    mapping_type t = get_mapping_type(options.map_file.c_str());

    if(t & PERM_FILE)
    {
        SequenceMappingParser parser(options.map_file.c_str());
        std::vector<unsigned int> read_snp_cand;
        read_snp_cand.resize(read_length);

        int total                         = 0;
        int skipped_cutoff                = 0;
        int skipped_unique                = 0;
        
        /* Iterate over the mapping file. */
        SeqRead r;
        while(parser.next_read(r))
        {
            total++;

            if( options.unique && r.map_locs.size() > 1)
            {
                skipped_unique++;
                continue;
            }

            if( r.map_locs.size() > options.max_map_loci )
            {
                skipped_cutoff++;
                continue;
            }

            for(unsigned int j = 0; j < r.map_locs.size(); j++)
            {
                MapLoci m = r.map_locs[j];

                //std::cout << "m.id = " << m.id << " m.pos = " << m.pos << std::endl;
                if(totl.find(m.id) == totl.end())
                {
                    if(options.ignore_missing)
                        continue;
                    else
                        error("chromosome " + m.id + " does not exist in the fasta index file. (did you forget --ignore_missing?)");
                }

                for(unsigned int k = 0; k < r.seq.size(); k++)
                    read_snp_cand[k] = 0;
                

                // add read to the totl array.
                for(unsigned int i = 0; i < r.seq.size(); i++)
                    totl[m.id][m.pos+i] += 1.0 / r.map_locs.size();

                // mark the substitution locations on the genome. 
                for(unsigned int k = 0; k < r.seq.size(); k++)
                {
                    if( r.seq[k] != m.seq[k] )
                    {
                        if(m.strand == '+')
                            subs[m.id][m.pos + k] += 1.0 / r.map_locs.size();
                        else // m.strand == '-'
                            subs[m.id][m.pos + r.seq.size() - k - 1] += 1.0 / r.map_locs.size();
                    }
                }
            }
        }
    }
    else if(t & SAM_FILE)
    {
        std::map<std::string, std::string> seqs;
        parse_fasta(options.fasta_file.c_str(), seqs);
        SamSequenceMappingParser parser(options.map_file.c_str(), seqs, t);
        std::vector<unsigned int> read_snp_cand;
        read_snp_cand.resize(read_length);

        int total                         = 0;
        int skipped_cutoff                = 0;
        int skipped_unique                = 0;
        
        /* Iterate over the mapping file. */
        SeqRead r;
        while(parser.next_read(r))
        {
            total++;

            if( options.unique && r.map_locs.size() > 1)
            {
                skipped_unique++;
                continue;
            }

            if( r.map_locs.size() > options.max_map_loci )
            {
                skipped_cutoff++;
                continue;
            }

            for(unsigned int j = 0; j < r.map_locs.size(); j++)
            {
                MapLoci m = r.map_locs[j];

                //std::cout << "m.id = " << m.id << " m.pos = " << m.pos << std::endl;
                if(totl.find(m.id) == totl.end())
                {
                    if(options.ignore_missing)
                        continue;
                    else
                        error("chromosome " + m.id + " does not exist in the fasta index file. (did you forget --ignore_missing?)");
                }
                
                for(unsigned int k = 0; k < r.seq.size(); k++)
                    read_snp_cand[k] = 0;
                
                // add read to the totl array.
                for(unsigned int i = 0; i < r.seq.size(); i++)
                    totl[m.id][m.pos+i] += 1.0 / r.map_locs.size();

                // mark the substitution locations on the genome. 
                for(unsigned int k = 0; k < r.seq.size(); k++)
                {
                    if( r.seq[k] != m.seq[k] )
                    {
                        if(m.strand == '+')
                            subs[m.id][m.pos + k] += 1.0 / r.map_locs.size();
                        else // m.strand == '-'
                            subs[m.id][m.pos + r.seq.size() - k - 1] += 1.0 / r.map_locs.size();
                    }
                }
            }
        }
    }
    else
        error("Bad Sequence Mapping Parser Type.");

}

void find_seq_snp_cands(prog_options &options, cand_blocks_t &cand_blocks, cand_vector_t &snps, map_str_vafl_t &subs, map_str_vafl_t &totl)
{
    // We store the SNPs in two ways for fast lookup.  An array the size of the
    // genome, so we can quickly determine with a lookup whether the size is
    // covered with a SNP.  And a list of each snp block, for storing the information about the SNPs.

    map_str_vafl_t::iterator mi;

    // generate candidates from subs/totl.
    
    /* Our final set of snp candidates to return. */
    std::map<std::string, std::valarray<bool> > snp_cand;
    
    for(mi = subs.begin(); mi != subs.end(); ++mi)
    {
        std::string  s = mi->first;

        snp_cand[s].resize(mi->second.size());
        snp_cand[s] = false;

        for(unsigned int cti = 0; cti < options.cover_threshold.size(); cti++)
        {
            int coverage  = options.cover_threshold[cti].first;
            float threshold = options.cover_threshold[cti].second;

            snp_cand[s] |= (totl[s] >= (float)coverage) && (subs[s] / totl[s] > threshold);
        }
    }
        
    // push dummy block into pos = 0, so that we do not have to do -1
    // adjustment later to access the snps in the array.
    std::string fake_chr("FAKE SNP");
    std::valarray<float> _subs, _totl;
    snps.push_back(SnpRecord( -1, fake_chr, -1, -1, _subs, _totl) );
 
    // TODO: is there a better way to do this with valarray.cshift()?

    unsigned int snp_id = 1;

    // Find beginning and end of snp blocks, then walk over sequence to mark blocks.
    for(mi = subs.begin(); mi != subs.end(); ++mi)
    {
        std::string chr = mi->first;
        cand_blocks[chr].resize(mi->second.size(), 0);

        for(unsigned int i = 0; i < mi->second.size(); i++)
            if(snp_cand[chr][i])
            {
                _subs.resize(1); _totl.resize(1);
                _subs = subs[chr][std::slice(i, 1, 1)];
                _totl = totl[chr][std::slice(i, 1, 1)];

                cand_blocks[chr][i] = snp_id;
                snps.push_back(SnpRecord(snp_id, chr, i, i+1, _subs, _totl));
                snp_id++;
            }
    }
    if(snp_id == 1)
    { 
        std::cerr << "comb did not generate any snp candidates!" << std::endl;
        std::cerr << "Consider modifying the cover/threshold parameters." << std::endl;
    }
}
